package com.rocpenghua.service;

import com.google.gson.Gson;
import com.rocpenghua.config.PluginSettings;
import com.google.gson.JsonObject;
import com.google.gson.JsonParser;
import okhttp3.*;

import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;

public class DeepSeekAIService implements AIService {
    private static final MediaType JSON = MediaType.get("application/json; charset=utf-8");
    private final OkHttpClient client;
    private final PluginSettings settings;

    public DeepSeekAIService() {
        this.settings = PluginSettings.getInstance();
        this.client = new OkHttpClient.Builder()
                .connectTimeout(30, TimeUnit.SECONDS)
                .readTimeout(30, TimeUnit.SECONDS)
                .writeTimeout(30, TimeUnit.SECONDS)
                .build();
    }

    @Override
    public String explainCode(String code) {
        String prompt = "Please explain the following code in detail:\n\n" + code;
        return callAPI(prompt);
    }

    @Override
    public String optimizeCode(String code) {
        String prompt = "Please optimize the following code and explain the changes:\n\n" + code;
        return callAPI(prompt);
    }

    @Override
    public String generateComments(String code) {
        String prompt = "Please add detailed comments to the following code:\n\n" + code;
        return callAPI(prompt);
    }

    @Override
    public String chat(String message) {
        return callAPI(message);
    }

    private String callAPI(String prompt) {
        if (settings.getApiKey().isEmpty()) {
            return "API key is not configured. Please set it in the plugin settings.";
        }


        Map<String,Object> requestBody = new HashMap<>();
        requestBody.put("model", "deepseek-chat");
        List<Map<String,Object>> messages = new ArrayList<>();
        Map<String,Object> message = new HashMap<>();
        message.put("role","user");
        message.put("content",prompt);
        messages.add(message);
        requestBody.put("messages", messages);
        requestBody.put("max_tokens", 1000);
        requestBody.put("temperature", 0.7);
        Gson gson = new Gson();
        String json = gson.toJson(requestBody);
        Request request = new Request.Builder()
                .url(settings.getBaseUrl() + "/completions")
                .header("Authorization", "Bearer " + settings.getApiKey())
                .post(RequestBody.create(json, JSON))
                .build();

        try (Response response = client.newCall(request).execute()) {
            if (!response.isSuccessful()) {
                return "Error: " + response.code() + " - " + response.message() + response.body();
            }

            String responseBody = response.body().string();
            JsonObject jsonResponse = JsonParser.parseString(responseBody).getAsJsonObject();
            return jsonResponse.getAsJsonArray("choices")
                    .get(0).getAsJsonObject()
                    .get("message").getAsJsonObject()
                    .get("content").getAsString();
        } catch (IOException e) {
            return "Error calling API: " + e.getMessage();
        }
    }
}